{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pretty printing has been turned ON\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import sys\n",
    "from six import StringIO, b\n",
    "from pprint import PrettyPrinter\n",
    "%pprint\n",
    "\n",
    "from gym import utils\n",
    "from gym.envs.toy_text import discrete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "pp = PrettyPrinter(indent=2)\n",
    "\n",
    "UP = 0\n",
    "RIGHT = 1\n",
    "DOWN = 2\n",
    "LEFT = 3\n",
    "\n",
    "MAPS = {'4x4':[\"SOOO\",\"OXOX\",\"OOOX\",\"XOOG\"]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class GridworldEnv(discrete.DiscreteEnv):\n",
    "    \"\"\"\n",
    "    FrozenLakeEnv1 is a copy environment from GYM toy_text FrozenLake-01\n",
    "\n",
    "    You are an agent on an 4x4 grid and your goal is to reach the terminal\n",
    "    state at the bottom right corner.\n",
    "    \n",
    "    For example, a 4x4 grid looks as follows:\n",
    "    \n",
    "    S  O  O  O\n",
    "    O  X  O  X\n",
    "    O  O  O  X\n",
    "    X  O  O  G\n",
    "    \n",
    "    S : starting point, safe\n",
    "    O : frozen surface, safe\n",
    "    X : hole, fall to your doom\n",
    "    G : goal, where the frisbee is located\n",
    "    \n",
    "    The episode ends when you reach the goal or fall in a hole.\n",
    "    You receive a reward of 1 if you reach the goal, and zero otherwise.\n",
    "    \n",
    "    You can take actions in each direction (UP=0, RIGHT=1, DOWN=2, LEFT=3).\n",
    "    Actions going off the edge leave you in your current state.\n",
    "    \"\"\"\n",
    "    metadata = {'render.modes': ['human', 'ansi']}\n",
    "    \n",
    "    def __init__(self, desc=None, map_name='4x4'):\n",
    "        self.desc = desc = np.asarray(MAPS[map_name], dtype='c')\n",
    "        self.nrow, self.ncol = nrow, ncol = desc.shape\n",
    "        self.shape = desc.shape\n",
    "        \n",
    "        nA = 4                    # 动作集个数\n",
    "        nS = np.prod(desc.shape)  # 状态集个数\n",
    "\n",
    "        MAX_Y = desc.shape[0]\n",
    "        MAX_X = desc.shape[1]\n",
    "\n",
    "        # initial state distribution [ 1.  0.  0.  ...] \n",
    "        isd = np.array(desc == b'S').astype('float64').ravel()\n",
    "        isd /= isd.sum()\n",
    "        \n",
    "        P = {}          \n",
    "        state_grid = np.arange(nS).reshape(self.shape)\n",
    "        it = np.nditer(state_grid, flags=['multi_index'])\n",
    "        \n",
    "        while not it.finished:\n",
    "            s = it.iterindex\n",
    "            y, x = it.multi_index\n",
    "\n",
    "            # P[s][a] == [(probability, nextstate, reward, done), ...]\n",
    "            P[s] = {a : [] for a in range(nA)}\n",
    "\n",
    "            s_letter = desc[y][x]\n",
    "            is_done = lambda letter: letter in b'GX'\n",
    "            reward = 0.0 if s_letter in b'G' else -1.0\n",
    "            \n",
    "            if is_done(s_letter):\n",
    "                P[s][UP] = [(1.0, s, reward, True)]\n",
    "                P[s][RIGHT] = [(1.0, s, reward, True)]\n",
    "                P[s][DOWN] = [(1.0, s, reward, True)]\n",
    "                P[s][LEFT] = [(1.0, s, reward, True)]\n",
    "            else:\n",
    "                ns_up = s if y == 0 else s - MAX_X\n",
    "                ns_right = s if x == (MAX_X - 1) else s + 1\n",
    "                ns_down = s if y == (MAX_Y - 1) else s + MAX_X\n",
    "                ns_left = s if x == 0 else s - 1\n",
    "\n",
    "                sl_up = desc[ns_up//MAX_Y][ns_up%MAX_X]\n",
    "                sl_right = desc[ns_right//MAX_Y][ns_right%MAX_X]\n",
    "                sl_down = desc[ns_down//MAX_Y][ns_down%MAX_X]\n",
    "                sl_left = desc[ns_left//MAX_Y][ns_left%MAX_X]\n",
    "                \n",
    "                P[s][UP] = [(1.0, ns_up, reward, is_done(sl_up))]\n",
    "                P[s][RIGHT] = [(1.0, ns_right, reward, is_done(sl_right))]\n",
    "                P[s][DOWN] = [(1.0, ns_down, reward, is_done(sl_down))]\n",
    "                P[s][LEFT] = [(1.0, ns_left, reward, is_done(sl_left))]\n",
    "                \n",
    "            it.iternext()\n",
    "                \n",
    "        self.P = P\n",
    "        \n",
    "        super(GridworldEnv, self).__init__(nS, nA, P, isd)\n",
    "\n",
    "    def _render(self, mode='human', close=False):\n",
    "        if close: # 初始化环境Environment的时候不显示\n",
    "            return\n",
    "        \n",
    "        outfile = StringIO() if mode == 'ansi' else sys.stdout\n",
    "\n",
    "        desc = self.desc.tolist()\n",
    "        desc = [[c.decode('utf-8') for c in line] for line in desc]\n",
    "        \n",
    "        state_grid = np.arange(self.nS).reshape(self.shape)\n",
    "        it = np.nditer(state_grid, flags=['multi_index'])\n",
    "        \n",
    "        while not it.finished:\n",
    "            s = it.iterindex\n",
    "            y, x = it.multi_index\n",
    "        \n",
    "            # 对于当前状态用红色标注\n",
    "            if self.s == s:\n",
    "                desc[y][x] = utils.colorize(desc[y][x], \"red\", highlight=True)\n",
    "            \n",
    "            it.iternext()\n",
    "       \n",
    "        outfile.write(\"\\n\".join(' '.join(line) for line in desc)+\"\\n\")\n",
    "\n",
    "        if mode != 'human':\n",
    "            return outfile\n",
    "        \n",
    "env = GridworldEnv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[41mS\u001b[0m O O O\n",
      "O X O X\n",
      "O O O X\n",
      "X O O G\n",
      "action:0(Up)\n",
      "done:False, observation:0, reward:-1.0\n",
      "\u001b[41mS\u001b[0m O O O\n",
      "O X O X\n",
      "O O O X\n",
      "X O O G\n",
      "action:1(Right)\n",
      "done:False, observation:1, reward:-1.0\n",
      "S \u001b[41mO\u001b[0m O O\n",
      "O X O X\n",
      "O O O X\n",
      "X O O G\n",
      "action:2(Down)\n",
      "done:True, observation:5, reward:-1.0\n",
      "{ 0: { 0: [(1.0, 0, -1.0, False)],\n",
      "       1: [(1.0, 1, -1.0, False)],\n",
      "       2: [(1.0, 4, -1.0, False)],\n",
      "       3: [(1.0, 0, -1.0, False)]},\n",
      "  1: { 0: [(1.0, 1, -1.0, False)],\n",
      "       1: [(1.0, 2, -1.0, False)],\n",
      "       2: [(1.0, 5, -1.0, True)],\n",
      "       3: [(1.0, 0, -1.0, False)]},\n",
      "  2: { 0: [(1.0, 2, -1.0, False)],\n",
      "       1: [(1.0, 3, -1.0, False)],\n",
      "       2: [(1.0, 6, -1.0, False)],\n",
      "       3: [(1.0, 1, -1.0, False)]},\n",
      "  3: { 0: [(1.0, 3, -1.0, False)],\n",
      "       1: [(1.0, 3, -1.0, False)],\n",
      "       2: [(1.0, 7, -1.0, True)],\n",
      "       3: [(1.0, 2, -1.0, False)]},\n",
      "  4: { 0: [(1.0, 0, -1.0, False)],\n",
      "       1: [(1.0, 5, -1.0, True)],\n",
      "       2: [(1.0, 8, -1.0, False)],\n",
      "       3: [(1.0, 4, -1.0, False)]},\n",
      "  5: { 0: [(1.0, 5, -1.0, True)],\n",
      "       1: [(1.0, 5, -1.0, True)],\n",
      "       2: [(1.0, 5, -1.0, True)],\n",
      "       3: [(1.0, 5, -1.0, True)]},\n",
      "  6: { 0: [(1.0, 2, -1.0, False)],\n",
      "       1: [(1.0, 7, -1.0, True)],\n",
      "       2: [(1.0, 10, -1.0, False)],\n",
      "       3: [(1.0, 5, -1.0, True)]},\n",
      "  7: { 0: [(1.0, 7, -1.0, True)],\n",
      "       1: [(1.0, 7, -1.0, True)],\n",
      "       2: [(1.0, 7, -1.0, True)],\n",
      "       3: [(1.0, 7, -1.0, True)]},\n",
      "  8: { 0: [(1.0, 4, -1.0, False)],\n",
      "       1: [(1.0, 9, -1.0, False)],\n",
      "       2: [(1.0, 12, -1.0, True)],\n",
      "       3: [(1.0, 8, -1.0, False)]},\n",
      "  9: { 0: [(1.0, 5, -1.0, True)],\n",
      "       1: [(1.0, 10, -1.0, False)],\n",
      "       2: [(1.0, 13, -1.0, False)],\n",
      "       3: [(1.0, 8, -1.0, False)]},\n",
      "  10: { 0: [(1.0, 6, -1.0, False)],\n",
      "        1: [(1.0, 11, -1.0, True)],\n",
      "        2: [(1.0, 14, -1.0, False)],\n",
      "        3: [(1.0, 9, -1.0, False)]},\n",
      "  11: { 0: [(1.0, 11, -1.0, True)],\n",
      "        1: [(1.0, 11, -1.0, True)],\n",
      "        2: [(1.0, 11, -1.0, True)],\n",
      "        3: [(1.0, 11, -1.0, True)]},\n",
      "  12: { 0: [(1.0, 12, -1.0, True)],\n",
      "        1: [(1.0, 12, -1.0, True)],\n",
      "        2: [(1.0, 12, -1.0, True)],\n",
      "        3: [(1.0, 12, -1.0, True)]},\n",
      "  13: { 0: [(1.0, 9, -1.0, False)],\n",
      "        1: [(1.0, 14, -1.0, False)],\n",
      "        2: [(1.0, 13, -1.0, False)],\n",
      "        3: [(1.0, 12, -1.0, True)]},\n",
      "  14: { 0: [(1.0, 10, -1.0, False)],\n",
      "        1: [(1.0, 15, -1.0, True)],\n",
      "        2: [(1.0, 14, -1.0, False)],\n",
      "        3: [(1.0, 13, -1.0, False)]},\n",
      "  15: { 0: [(1.0, 15, 0.0, True)],\n",
      "        1: [(1.0, 15, 0.0, True)],\n",
      "        2: [(1.0, 15, 0.0, True)],\n",
      "        3: [(1.0, 15, 0.0, True)]}}\n",
      "Episode finished after 3 timesteps\n"
     ]
    }
   ],
   "source": [
    "observation = env.reset()\n",
    "for _ in range(5):\n",
    "    env.render()\n",
    "    action = env.action_space.sample()\n",
    "    observation, reward, done, info = env.step(action)\n",
    "    print(\"action:{}({})\".format(action, [\"Up\",\"Right\",\"Down\",\"Left\"][action]))\n",
    "    print(\"done:{}, observation:{}, reward:{}\".format(done, observation, reward))\n",
    "    if done:\n",
    "        pp.pprint(env.P)\n",
    "        print(\"Episode finished after {} timesteps\".format(_+1))\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
